from paddlenlp import Taskflow


'''
NLP信息抽取，企业名
百度自研模型，通用信息抽取 UIE(Universal Information Extraction)
'''


def entity_extraction(text, entity):
    schema = [entity]    # Define the schema for entity extraction
    # 信息抽取应用使用ERNIE 3.0轻量级模型作为预训练模型，同时在大量信息抽取数据上进行了二次预训练，从而让模型适配固定prompt。
    # 百度文心ERNIE 3.0大语言模型(Large Language Model, LLM)擅长中文
    # 模型：uie-base (默认)
    # 结构：12-layers, 768-hidden, 12-heads
    # 语言：中文
    ie = Taskflow('information_extraction', schema=schema, model='uie-base')
    data = ie(text)
    # 创建一个空列表来存储提取的文本
    extracted_texts = []

    # 遍历列表中的每个字典
    for item in data:
        # 遍历字典中的'企业'键对应的列表
        for enterprise in item['企业']:
            # 提取'text'键对应的值，并添加到extracted_texts列表中
            extracted_texts.append(enterprise['text'])  # 提取的文本
    unique_list = list(dict.fromkeys(extracted_texts))  # 去重复
    return unique_list


text = '''
《股东来了》（2024）由中证中小投资者服务中心（简称“投资者服务中心”）联合上海、江苏、西藏、深圳四地证监局携手腾讯集团、芒果超媒共同举办。

本届活动分为网络初赛和电视决赛两个阶段。网络初赛采用线上答题形式，答题平台于5月15日8点正式面向公众开放，共持续8周。通过“投资者成长营”带投资者领略“读心魔法屋”“个人闯关赛”“猜词消消乐”“PK竞技场”四大游戏玩法！比赛拼脑力，答题学知识！期待投资者通过轻松愉快的方式进一步学习权益知识等相关内容。
'''
result = entity_extraction(text, entity='企业')
print(result)
# ['芒果超媒', '腾讯集团']


text = '''
IT市场年会由赛迪顾问股份有限公司主办，赛迪顾问还将重磅发布140个热点行业领域报告及数据，洞察未来发展趋势。此外，在赛迪顾问本届颁奖仪式中，英威腾电源的模块化不间断电源(UPS)荣获 “新一代信息技术创新产品”奖项荣誉。
'''
result = entity_extraction(text, entity='企业')
print(result)
# ['赛迪顾问', '赛迪顾问股份有限公司', '英威腾电源']


text = '''
论坛上，英威腾副总裁兼工业自动化事业群总裁杨林先生，代表英威腾致辞。英威腾项目销售总经理刘长志、国内功率半导体器件领域的领军企业——斯达半导体股份有限公司，副总经理许浩平、全国电容器行业优势企业——南通江海电容器股份有限公司，总裁丁继华、龙芯中科技术股份有限公司——高级总监黎旭安分别做了主题演讲。
'''
result = entity_extraction(text, entity='企业')
print(result)
# ['斯达半导体股份有限公司', '英威腾', '龙芯中科技术股份有限公司', '南通江海电容器股份有限公司']


'''
NLP关系抽取，是合作关系的企业
'''
import requests
import json
import re


def extract_first_number(text):
    match = re.search(r'\d+', text)
    if match:
        return match.group()
    else:
        return None


def get_access_token():
    """
    使用 API Key，Secret Key 获取access_token，替换下列示例中的应用API Key、应用Secret Key
    """
    url = "https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=API_KEY&client_secret=Secret_KEY"

    payload = ""
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload)

    return response.json().get("access_token")


def create_chat(dialogueRequest):
    # 调用API
    # ERNIE4.0，是百度公司发布的文心大模型版本。
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-4.0-8k-preview?access_token=" + get_access_token()

    payload = json.dumps({
        "messages": [
            {
                "role": "user",
                "content": dialogueRequest
            }
        ]
    })
    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload)

    # 提取result字段
    result = extract_first_number(json.loads(response.text)['result'])
    print(result)
    if result == "1":
        return "是合作关系"
    elif result == "2":
        return "不是合作关系"


dialogueRequest = '''
“《股东来了》（2024）由中证中小投资者服务中心（简称“投资者服务中心”）联合上海、江苏、西藏、深圳四地证监局携手腾讯集团、芒果超媒共同举办。
本届活动分为网络初赛和电视决赛两个阶段。网络初赛采用线上答题形式，答题平台于5月15日8点正式面向公众开放，共持续8周。通过“投资者成长营”带投资者领略“读心魔法屋”“个人闯关赛”“猜词消消乐”“PK竞技场”四大游戏玩法！比赛拼脑力，答题学知识！期待投资者通过轻松愉快的方式进一步学习权益知识等相关内容。”请问在这段话中，腾讯集团和芒果超媒是合作伙伴关系吗？如果是合作伙伴关系请返回1，不是合作伙伴关系请返回2。直接返回1或者2即可，不必给出其他理由等。
'''
result = create_chat(dialogueRequest)
print(result)
# 1
# 是合作关系


dialogueRequest = '''
“IT市场年会由赛迪顾问股份有限公司主办，赛迪顾问还将重磅发布140个热点行业领域报告及数据，洞察未来发展趋势。此外，在赛迪顾问本届颁奖仪式中，英威腾电源的模块化不间断电源(UPS)荣获 “新一代信息技术创新产品”奖项荣誉。”请问在这段话中，赛迪顾问股份有限公司和英威腾是合作伙伴关系吗？如果是合作伙伴关系请返回1，不是合作伙伴关系请返回2。直接返回1或者2即可，不必给出其他理由等。
'''
result = create_chat(dialogueRequest)
print(result)
# 2
# 不是合作关系


dialogueRequest = '''
“论坛上，英威腾副总裁兼工业自动化事业群总裁杨林先生，代表英威腾致辞。英威腾项目销售总经理刘长志、国内功率半导体器件领域的领军企业——斯达半导体股份有限公司，副总经理许浩平、全国电容器行业优势企业——南通江海电容器股份有限公司，总裁丁继华、龙芯中科技术股份有限公司——高级总监黎旭安分别做了主题演讲。”请问在这段话中，龙芯中科技术股份有限公司和英威腾是合作伙伴关系吗？如果是合作伙伴关系请返回1，不是合作伙伴关系请返回2。直接返回1或者2即可，不必给出其他理由等。
'''
result = create_chat(dialogueRequest)
print(result)
# 2
# 不是合作关系
